Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Modified truncate_normal with jax.random.truncated_normal #175

Merged
merged 7 commits into from
Jul 24, 2024

Conversation

mhliu0001
Copy link
Contributor

Previously the truncate_normal distribution was sampled from a normal distribution and then clipping the result. This does not yield a continuous distribution. A more reasonable distribution should be like: https://en.wikipedia.org/wiki/Truncated_normal_distribution

This should affect the microphysics simulation. Maybe the previous version is intentional, so I will let experts comment on which behavior is more correct.

@coveralls
Copy link

coveralls commented Jul 18, 2024

Pull Request Test Coverage Report for Build 9993829303

Details

  • 12 of 14 (85.71%) changed or added relevant lines in 1 file are covered.
  • 3 unchanged lines in 1 file lost coverage.
  • Overall coverage decreased (-0.1%) to 84.711%

Changes Missing Coverage Covered Lines Changed/Added Lines %
appletree/randgen.py 12 14 85.71%
Files with Coverage Reduction New Missed Lines %
appletree/randgen.py 3 71.03%
Totals Coverage Status
Change from base Build 9897876272: -0.1%
Covered Lines: 2377
Relevant Lines: 2806

💛 - Coveralls

@mhliu0001
Copy link
Contributor Author

jax-ml/jax#10951
When using jax.random.truncated_normal, we should be careful not to set vmin/vmax too extreme.

But I still think this should be better than using the previous clipping method. The comparison when simulating a truncated normal distribution with mean=0, std=1, vmin=-2, vmax=2 is shown below.
truncate_normal

Copy link
Collaborator

@zihaoxu98 zihaoxu98 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks! The PR looks good to me

@zihaoxu98 zihaoxu98 merged commit a74b86f into master Jul 24, 2024
7 checks passed
@zihaoxu98 zihaoxu98 deleted the truncate_normal branch July 24, 2024 15:49
@mhliu0001
Copy link
Contributor Author

We should revert this PR because NEST also uses the previous truncate_normal in quanta generation. Here NEST defines two types of truncate_normal, one is the non-continuous version (termed rand_gauss with zero_min == true), and the other is the continuous version defined by this PR (termed rand_zero_trunc_gauss). In the quanta generation part, NEST uses rand_gauss so the previous version is more correct.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants